package cn.zhl.tree.avl;

import cn.zhl.tree.brst.BinaryRotateSearchTree;

import java.util.Comparator;

/**
 * @program: demos
 * @description:
 * @author: 刘振华
 * @create: 2020-05-03 17:08
 **/
public class AVLTree<E> extends BinaryRotateSearchTree<E> {
	public AVLTree() {
	}

	public AVLTree(Comparator<E> comparator) {
		super(comparator);
	}

	@Override
	protected void afterAdd(Node node) {
		while ((node = node.parent) != null){
			AVLNode anode = (AVLNode) node;
			if(isBalance(anode)){
				//更新高度
				anode.updateHeight();
			}else {
				//回复平衡
				reBalance(anode);
				break;
			}
		}
	}

	@Override
	protected void afterRemove(Node node) {
		while ((node = node.parent) != null){
			AVLNode anode = (AVLNode) node;
			if(isBalance(anode)){
				//更新高度
				anode.updateHeight();
			}else {
				//回复平衡
				reBalance(anode);
			}
		}
	}

	//核心，恢复平衡
	private void reBalance(AVLNode node){
		Node first = node.tallerChild();
		Node second = ((AVLNode) first).tallerChild();
		if(first.isLeftChild()){
			if(second.isLeftChild()){//LL
				rotateRight(node);
			}else {//LR
				rotateLeft(first);
				rotateRight(node);
			}
		}else {
			if(second.isLeftChild()){//RL
				rotateRight(first);
				rotateLeft(node);
			}else {//RR
				rotateLeft(node);
			}
		}
	}

	@Override
	protected void afterRotate(Node grand, Node parent, Node child) {
		super.afterRotate(grand, parent, child);
		((AVLTree.AVLNode) grand).updateHeight();
		((AVLTree.AVLNode) parent).updateHeight();
	}

	private boolean isBalance(AVLNode node){
		return Math.abs(node.balanceFactor()) <= 1;
	}

	@Override
	protected AVLNode createNode(E e, Node parent) {
		return new AVLNode(e,parent);
	}

	private class AVLNode extends Node{
		int height = 1;
		public AVLNode(E element, Node parent) {
			super(element, parent);
		}

		public int balanceFactor(){
			int leftHeight = left == null ? 0 : ((AVLNode) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode) right).height;
			return leftHeight - rightHeight;
		}

		public void updateHeight(){
			int leftHeight = left == null ? 0 : ((AVLNode) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode) right).height;
			height = 1 + Math.max(leftHeight,rightHeight);
		}

		public Node tallerChild(){
			int leftHeight = left == null ? 0 : ((AVLNode) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode) right).height;
			if(leftHeight > rightHeight){
				return left;
			}else if(leftHeight < rightHeight){
				return right;
			}else {
				return isLeftChild() ? left : right;
			}
		}

	}
}
